import random
from random import randint

import gym
import numpy as np
import pygame
from cpgames.core import PingpongGame
from cpgames.core.games.pingpong.modules import Racket, Ball
from cpgames.core.games.pingpong.pingpong import Config
from cpgames.core.utils import QuitGame, InitPygame, PygameResourceLoader


class NoSound:
    def play(self):
        pass


class AIBall(Ball):

    def reset(self):
        super().reset()
        self.speed = random.randint(1, 5)

    def move(self, ball, racket_left, racket_right, hit_sound, goal_sound):
        scores = super(AIBall, self).move(ball, racket_left, racket_right, hit_sound, goal_sound)
        # self.speed = min(self.speed, 10)
        return scores


class AIPingPongGame(PingpongGame):

    def __init__(self, **kwargs):
        self.cfg = Config()
        # 设置属性
        self.config = Config
        # 初始化
        self.initialize()
        # 用户可以覆盖默认参数
        for key, value in kwargs.items():
            if hasattr(self, key): setattr(self, key, value)

    def initialize(self):
        self.screen = InitPygame(screensize=self.config.SCREENSIZE, title=self.config.TITLE, init_mixer=False)
        bgm_path = self.config.BGM_PATH if hasattr(self.config, 'BGM_PATH') else None
        font_paths_dict = self.config.FONT_PATHS_DICT if hasattr(self.config, 'FONT_PATHS_DICT') else None
        image_paths_dict = self.config.IMAGE_PATHS_DICT if hasattr(self.config, 'IMAGE_PATHS_DICT') else None
        # sound_paths_dict = self.config.SOUND_PATHS_DICT if hasattr(self.config, 'SOUND_PATHS_DICT') else None
        self.resource_loader = PygameResourceLoader(
            bgm_path=bgm_path,
            font_paths_dict=font_paths_dict,
            image_paths_dict=image_paths_dict,
        )

    def AIGamingInterface(self, resource_loader, cfg):
        # 初始化
        # hit_sound = resource_loader.sounds['hit']
        # goal_sound = resource_loader.sounds['goal']
        hit_sound = NoSound()
        goal_sound = NoSound()
        font = resource_loader.fonts['default50']
        # resource_loader.playbgm()
        # 开始界面
        # game_mode = 1
        # 游戏主循环
        # --左边球拍(ws控制, 仅双人模式时可控制)
        # score_left = 0
        racket_left = Racket(cfg.IMAGE_PATHS_DICT['racket'], 'LEFT', cfg)
        # --右边球拍(↑↓控制)
        # score_right = 0
        racket_right = Racket(cfg.IMAGE_PATHS_DICT['racket'], 'RIGHT', cfg)
        racket_right.speed = 10
        # --球
        ball = AIBall(cfg.IMAGE_PATHS_DICT['ball'], cfg)

        return racket_left, racket_right, ball, hit_sound, goal_sound, font


class PingPong2DEnv(gym.Env):
    def __init__(self):
        self.game = AIPingPongGame()
        self.score_left = 0
        self.score_right = 0
        self.clock = pygame.time.Clock()
        self.last_action = 2
        self.obs = None
        self._max_episode_steps = 10000
        self.observation_space = gym.spaces.Box(low=0, high=1, shape=(125, 125), dtype=np.float32)
        self.action_space = gym.spaces.Discrete(3)

    def reset(self, *args, **kwargs):
        racket_left, racket_right, ball, hit_sound, goal_sound, font = self.game.AIGamingInterface(
            self.game.resource_loader,
            self.game.cfg
        )
        self.racket_left = racket_left
        self.racket_right = racket_right

        # self.score_right = 0
        self.score_left = 0
        # self.score_right = randint(0, 9)
        self.score_right = 0
        self.ball = ball
        self.hit_sound = hit_sound
        self.goal_sound = goal_sound
        self.last_action = 2
        self.font = font
        # 返回状态, 用于初始化 [球速, 球位置, 球方向, 左拍位置, 右拍位置]
        one_frame = pygame.surfarray.array3d(self.game.screen)
        # （500, 500, 3）-> (3, 500, 500)
        one_frame = np.transpose(one_frame, (2, 0, 1))
        # 灰度化 0.3R + 0.59G + 0.11B
        one_frame = one_frame[0] * 0.3 + one_frame[1] * 0.59 + one_frame[2] * 0.11
        # 二值化
        one_frame[one_frame <= 100] = 0
        one_frame[one_frame > 100] = 1
        # 没隔一个像素取一个像素 500 -> 250 -> 125
        one_frame = one_frame[::2, ::2]
        one_frame = one_frame[::2, ::2]

        self.obs = one_frame
        return self.obs

    def step(self, action):
        # action = 0, 1, 2分别代表上，下，不动
        if action == 0:
            self.racket_right.move('UP')
        elif action == 1:
            self.racket_right.move('DOWN')
        elif action == 2:
            pass
        # 球运动
        scores = self.ball.move(self.ball, self.racket_left, self.racket_right, self.hit_sound, self.goal_sound)
        self.racket_left.automove(self.ball)
        self.score_left += scores[0]
        self.score_right += scores[1]
        done = False
        if self.score_left == 11 or self.score_right == 11:
            done = True
            # reward = -abs(self.racket_right.rect.centery - self.ball.rect.centery) * 0.1
        # 掉球惩罚
        racket_height = self.racket_right.rect.bottom - self.racket_right.rect.top
        if scores[0] > 0:
            # 距离中心越远,分数越低
            reward = -(abs(self.racket_right.rect.centery - self.ball.rect.centery) / racket_height) * 2 -5
        # elif scores[1] > 0:
        #     reward = 1
        # 打中奖励
        elif pygame.sprite.collide_rect(self.ball, self.racket_right):
            # 计算球拍高度
            # 距离中心越近，分数越高
            if abs(self.ball.rect.centery - self.racket_right.rect.centery) < racket_height / 4:
                reward = 5+(racket_height / 2 - abs(
                    self.ball.rect.centery - self.racket_right.rect.centery)) / racket_height / 5
            else:
                reward = 0
        else:
            reward = 0
        # if abs(self.racket_right.rect.centery - self.ball.rect.centery) > 20:
        #     reward += -0.001
        # if action == self.last_action:
        #     reward += 0.01
        # else:
        #     reward += -0.01
        self.last_action = action
        pygame.display.update()
        surf_array = pygame.surfarray.array3d(self.game.screen)
        # (500, 500, 3)转成（3，500， 500）的形式
        surf_array = np.transpose(surf_array, (2, 0, 1))
        # 灰度化 0.3R + 0.59G + 0.11B, 图片压缩到80*80
        surf_array = surf_array[0] * 0.3 + surf_array[1] * 0.59 + surf_array[2] * 0.11
        # 二值化 大于100的像素点为1，小于100的像素点为0
        surf_array[surf_array <= 100] = 0
        surf_array[surf_array > 100] = 1
        # 没隔一个像素取一个像素 500 -> 250 -> 125
        surf_array = surf_array[::2, ::2]
        surf_array = surf_array[::2, ::2]
        self.obs = surf_array

        # 返回状态, 奖励, 是否结束, 其他信息
        return self.obs, reward, done, {}

    def render(self, mode='human'):
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                QuitGame()
        self.game.screen.fill((41, 36, 33))
        # 玩家操作
        pressed_keys = pygame.key.get_pressed()
        if pressed_keys[pygame.K_UP]:
            self.racket_right.move('UP')
        elif pressed_keys[pygame.K_DOWN]:
            self.racket_right.move('DOWN')
        self.racket_left.automove(self.ball)
        # 球运动
        scores = self.ball.move(self.ball, self.racket_left, self.racket_right, self.hit_sound, self.goal_sound)
        self.score_left += scores[0]
        self.score_right += scores[1]
        # 显示
        # --分隔线
        pygame.draw.rect(self.game.screen, self.game.cfg.WHITE, (247, 0, 6, 500))
        # --球
        self.ball.draw(self.game.screen)
        # --拍
        self.racket_left.draw(self.game.screen)
        self.racket_right.draw(self.game.screen)
        # --得分
        self.game.screen.blit(self.font.render(str(self.score_left), False, self.game.cfg.WHITE), (150, 10))
        self.game.screen.blit(self.font.render(str(self.score_right), False, self.game.cfg.WHITE), (300, 10))
        pygame.display.update()
        # self.clock.tick(self.game.cfg.FPS_GAMING)


if __name__ == '__main__':
    pingpone_env = PingPong2DEnv(AIPingPongGame())
    pingpone_env.game.run()
    while True:
        state = pingpone_env.reset()
        done = False
        while not done:
            # action = random.randint(0, 3)
            # state, reward, done, _ = pingpone_env.step(action)
            pingpone_env.render()
            # print(state, reward, done)
